from Feature_extraction import Feature_extraction
import time
import warnings
warnings.filterwarnings('ignore')
import os
from tqdm import tqdm
from multiprocessing import Process
import numpy as np
import pandas as pd
from watchDog import start_monitoring
if __name__ == '__main__':

    start = time.time()
    print("========== CIC IoT feature extraction ==========")
    
    pcapfiles = start_monitoring()
    subfiles_size = 10 # MB
    split_directory = 'split_temp/'
    destination_directory = 'output/'
    converted_csv_files_directory = 'csv_files/'
    n_threads = 8
    
    address = "./"
    
        

    
    for i in range(len(pcapfiles)):
        lstart = time.time()
        pcap_file = pcapfiles[i]
        print(pcap_file)
        print(">>>> 1. splitting the .pcap file.")
        os.system('tcpdump -r '+ pcap_file +' -w ' + split_directory + 'split_temp -C ' + str(subfiles_size))
        subfiles = os.listdir(split_directory)
        print(">>>> 2. Converting (sub) .pcap files to .csv files.")
        processes = []
        errors = 0
        
        subfiles_threadlist = np.array_split(subfiles, (len(subfiles)/n_threads)+1)
        for f_list in tqdm(subfiles_threadlist):
            n_processes = min(len(f_list), n_threads)
            assert n_threads >= n_processes
            assert n_threads >= len(f_list)
            processes = []
            for i in range(n_processes):
                fe = Feature_extraction()
                f = f_list[i]
                subpcap_file = split_directory + f
                p = Process(target=fe.pcap_evaluation, args=(subpcap_file,destination_directory + f.split('.')[0]))
                p.start()
                processes.append(p)
            for p in processes:
                p.join()
        assert len(subfiles)==len(os.listdir(destination_directory))
        print(">>>> 3. Removing (sub) .pcap files.")
        for sf in subfiles:
            os.remove(split_directory + sf)

        print(">>>> 4. Merging (sub) .csv files (summary).")

        # 1. 检查目标目录是否存在
        if not os.path.exists(destination_directory):
            print(f"Error: Directory {destination_directory} does not exist!")
            exit(1)

        # 2. 获取CSV文件列表并检查是否为空
        csv_subfiles = [f for f in os.listdir(destination_directory) if f.endswith('.csv')]
        if not csv_subfiles:
            print(f"Error: No CSV files found in {destination_directory}")
            exit(1)

        print(f"Found {len(csv_subfiles)} CSV files to merge")

        # 3. 初始化合并参数
        output_csv = f"{pcap_file}.csv"
        mode = 'w'
        merged_count = 0

        # 4. 合并处理
        for f in tqdm(csv_subfiles, desc="Merging CSVs"):
            try:
                file_path = os.path.join(destination_directory, f)
                
                # 检查文件是否为空
                if os.path.getsize(file_path) == 0:
                    print(f"Warning: {f} is empty, skipping")
                    continue
                    
                # 读取CSV文件
                d = pd.read_csv(file_path)
                
                # 第一次写入表头，后续追加
                write_header = (mode == 'w')
                d.to_csv(
                    output_csv,
                    header=write_header,
                    index=False,
                    mode=mode
                )
                mode = 'a'
                merged_count += 1
                
            except pd.errors.EmptyDataError:
                print(f"Warning: {f} is empty or corrupt, skipping")
            except Exception as e:
                print(f"Error processing {f}: {str(e)}")
                continue

        # 5. 验证结果
        if merged_count == 0:
            print("Error: No files were successfully merged")
            exit(1)

        if os.path.exists(output_csv):
            # 检查合并后的文件行数
            try:
                merged_df = pd.read_csv(output_csv)
                print(f"Successfully merged {merged_count} files")
                print(f"Merged file contains {len(merged_df)} rows")
                print(f"Merged file saved to: {output_csv}")
            except Exception as e:
                print(f"Error verifying merged file: {str(e)}")
        else:
            print("Error: Failed to generate merged CSV file")

        print(">>>> 5. Removing (sub) .csv files.")
        #for cf in tqdm(csv_subfiles):
            #os.remove(destination_directory + cf)
        print(f'done! ({pcap_file})(' + str(round(time.time()-lstart, 2))+ 's),  total_errors= '+str(errors))
        
    end = time.time()
    print(f'Elapsed Time = {(end-start)}s')
